import torch

def MSE(pred, gt):
    return torch.mean(torch.square(pred - gt))

def L1Loss(pred, gt):
    return torch.mean(torch.abs(pred-gt))

def MaxL1Loss4D(pred, gt):
    return torch.max(torch.mean(torch.abs(pred-gt), axis=[1,2,3]))

def MaxL1Loss2D(pred, gt):
    return torch.max(torch.abs(pred-gt))


def WeightsNormLoss(model, device):
    t_zero = torch.FloatTensor([0.0]).to(device)
    weight_norms = []
    for p in model.parameters():
        weight_norms += [torch.mean(torch.abs(p))]
    min_norm = torch.min(torch.stack(weight_norms))
    max_norm = torch.max(torch.stack(weight_norms))
    norm_loss = 1000*torch.max(0.001-min_norm, t_zero) + torch.max(max_norm - 1000.0, t_zero)
    return norm_loss
